/*
 * HapMatrices.cpp
 *
 *  Created on: Apr 12, 2012
 *      Author: pickrell
 */

#include "HapMatrices.h"

HapMatrices::HapMatrices(string infile, string indfile, string mixpop, string sourcepop){
	int seed =  unsigned ( time (NULL) );
	srand (seed);
	//cout <<"here\n";
	get_ind2keep(indfile, mixpop, sourcepop);
	//cout << "here2\n";
	read_hapmatrices(infile);
	cout << "Number of admixed individuals: "<< allmixinds.size() << "\n";
	cout <<"Number of source individuals: "<< sourceinds.size() << "\n";
	random_shuffle(allmixinds.begin(), allmixinds.end());
	for (int i = 0; i < allmixinds.size(); i++){
		if (i < allmixinds.size()/2 ) mixinds1.insert(allmixinds[i]);
		else (mixinds2.insert(allmixinds[i]));
	}
	MAXDIST = 0.4;
}

void HapMatrices::read_hapmatrices(string infile){
	ifstream in(infile.c_str());
	vector<string> line;
	struct stat stFileInfo;
	int intStat;
	string st, buf;
	intStat = stat(infile.c_str(), &stFileInfo);
	if (intStat !=0){
		std::cerr<< "ERROR: cannot open file " << infile << "\n";
		exit(1);
	}

	while(getline(in, st)){
		buf.clear();
		stringstream ss(st);
		line.clear();
		while (ss>> buf){
			line.push_back(buf);
		}
		string hapfile = line[0];
		string mapfile = line[1];
		cout << "Reading "<< hapfile << "\n";
		HapMatrix tmp(hapfile, mapfile, ind2keep);
		d.push_back(tmp);
	}

	//remove individuals in not in the data but in the population file
	map<string, bool> found = d[0].ind2found;
	set<string> ind2keep_bk;
	for (set<string>::iterator it = ind2keep.begin(); it != ind2keep.end(); it++){
		//cout << "Looking for "<< *it << "\n";
		if (found.find(*it) == found.end()){
			cerr << "WARNING: " << *it << " not found\n";
		}
		else{
			ind2keep_bk.insert(*it);
		}
	}
	ind2keep.clear();
	for (set<string>::iterator it = ind2keep_bk.begin(); it != ind2keep_bk.end(); it++) ind2keep.insert(*it);
	vector<string> allmixinds_bk;
	set<string> sourceinds_bk;
	for (vector<string>::iterator  it = allmixinds.begin(); it != allmixinds.end(); it++){
		if (ind2keep.find(*it) != ind2keep.end()) allmixinds_bk.push_back(*it);
	}
	for (set<string>::iterator it = sourceinds.begin(); it != sourceinds.end(); it++){
		if (ind2keep.find(*it) != ind2keep.end()) sourceinds_bk.insert(*it);
	}
	allmixinds.clear();
	sourceinds.clear();
	for (vector<string>::iterator  it = allmixinds_bk.begin(); it != allmixinds_bk.end(); it++) {
		//cout << "Adding "<< *it << "\n";
		allmixinds.push_back(*it);
	}

	for (set<string>::iterator it = sourceinds_bk.begin(); it != sourceinds_bk.end(); it++) {
		//cout << "Adding2 " <<*it << "\n";
		sourceinds.insert(*it);
	}
}


void HapMatrices::get_ind2keep(string popfile, string mixpop, string sourcepop){
	ind2keep.clear();
	ind2pop.clear();
	allmixinds.clear();
	sourceinds.clear();
	ifstream in(popfile.c_str());
	vector<string> line;
	struct stat stFileInfo;
	int intStat;
	string st, buf;
	intStat = stat(popfile.c_str(), &stFileInfo);
	if (intStat !=0){
		std::cerr<< "ERROR: cannot open file " << popfile << "\n";
		exit(1);
	}
	while(getline(in, st)){
		buf.clear();
		stringstream ss(st);
		line.clear();
		while (ss>> buf){
			line.push_back(buf);
		}
		string id = line[0];
		string popid = line[2];
		//cout << id << " "<< popid << "\n";
		if (popid == mixpop || popid == sourcepop){
			ind2keep.insert(id);
			ind2pop.insert(make_pair(id, popid));
			if (popid == mixpop) allmixinds.push_back(id);
			else if (popid == sourcepop) sourceinds.insert(id);
		}
	}
}

void HapMatrices::print(){
	for (vector<HapMatrix>::iterator it = d.begin(); it != d.end(); it++){
		it->print();
	}
}

int HapMatrices::set_weights(){
	int toreturn = 0;
	float f2_tmp = 0;
	float hzy_mix_tmp = 0;
	float hzy_source_tmp = 0;
	float c_mix = 0;
	float c_source = 0;
	for (vector<HapMatrix>::iterator it = d.begin(); it != d.end(); it++){
		toreturn += it->set_weights(mixinds1, sourceinds, mixinds2, &f2_tmp, &hzy_mix_tmp, &hzy_source_tmp, &c_mix, &c_source);
	}
	f2_uncor = f2_tmp / (double) toreturn;
	hzy_mix = hzy_mix_tmp / (double) toreturn;
	hzy_source = hzy_source_tmp / (double) toreturn;
	mean_nind_mix = c_mix / (double) toreturn;
	mean_nind_source = c_source / (double) toreturn;
	cor_mix = hzy_mix / (4* mean_nind_mix);
	cor_source = hzy_source / (4* mean_nind_source);
	f2 = f2_uncor - cor_mix - cor_source;

	//cout << f2 << " "<< f2_uncor << " "<< hzy_mix << " "<< hzy_source << " "<< mean_nind_mix << " "<< mean_nind_source << " "<< cor_mix << " "<< cor_source << "\n";
	beta_cor = (f2_uncor*f2_uncor) / (f2*f2);
	//cout << beta_cor << "\n";
	return toreturn;
}

void HapMatrices::print_d(){
	for (vector<HapMatrix>::iterator it = d.begin(); it != d.end(); it++){
		it->print_all_d(mixinds2, MAXDIST);
	}
}

vector<pair<vector<float>, vector<float> > > HapMatrices::bin_d(double by){
	//cout << "here\n"; cout.flush();
	vector<pair<vector<float>, vector<float> > > toreturn;
	for (double i  = 0; i < MAXDIST; i+= by){
		pair<vector<float>, vector<float> > tmp;
		toreturn.push_back(tmp);
	}
	for (vector<HapMatrix>::iterator it = d.begin(); it != d.end(); it++){
		vector<pair<vector<float>, vector<float> > > tmp = it->bin_d(mixinds2, 0, MAXDIST, by);
		//cout << "got binned\n";
		for (int i = 0; i < tmp.size(); i++){
			for (int j = 0; j < tmp[i].first.size(); j++){
				toreturn[i].first.push_back(tmp[i].first.at(j));
				toreturn[i].second.push_back(tmp[i].second.at(j));
			}
		}
	}
	return toreturn;
}


vector<vector<float> > HapMatrices::bin_d_lowmem(double by){
	//cout << "here\n"; cout.flush();
	vector<vector<float> > toreturn;
	for (double i  = 0; i < MAXDIST; i+= by){
		vector<float> tmp (5, 0.0);
		toreturn.push_back(tmp);
	}
	for (vector<HapMatrix>::iterator it = d.begin(); it != d.end(); it++){
		vector<vector<float> > tmp = it->bin_d_lowmem(mixinds2, 0, MAXDIST, by);
		cout << "Done with "<< it->nsnp << "\n";
		for (int i = 0; i < tmp.size(); i++){
			for (int j = 0; j < tmp[i].size(); j++){
				toreturn[i][j] += tmp[i][j];
			}
		}
	}
	return toreturn;
}


vector<vector<float> > HapMatrices::bin_d_freq_lowmem(double by){
	//cout << "here\n"; cout.flush();
	vector<vector<float> > toreturn;
	for (double i  = 0; i < MAXDIST; i+= by){
		vector<float> tmp (5, 0.0);
		toreturn.push_back(tmp);
	}
	for (vector<HapMatrix>::iterator it = d.begin(); it != d.end(); it++){
		vector<vector<float> > tmp = it->bin_d_freq_lowmem(mixinds2, 0, MAXDIST, by);
		cout << "Done with "<< it->nsnp << "\n";
		for (int i = 0; i < tmp.size(); i++){
			for (int j = 0; j < tmp[i].size(); j++){
				toreturn[i][j] += tmp[i][j];
			}
		}
	}
	return toreturn;
}



vector<vector<float> > HapMatrices::bin_d_freq_cor_lowmem(double by, vector<pair<float, float> > cor){
	//cout << "here\n"; cout.flush();
	vector<vector<float> > toreturn;
	for (double i  = 0; i < MAXDIST; i+= by){
		vector<float> tmp (5, 0.0);
		toreturn.push_back(tmp);
	}
	for (vector<HapMatrix>::iterator it = d.begin(); it != d.end(); it++){
		vector<vector<float> > tmp = it->bin_d_freq_cor_lowmem(mixinds2, cor, 0, MAXDIST, by);
		cout << "Done with "<< it->nsnp << "\n";
		for (int i = 0; i < tmp.size(); i++){
			for (int j = 0; j < tmp[i].size(); j++){
				toreturn[i][j] += tmp[i][j];
			}
		}
	}
	return toreturn;
}

double HapMatrices::get_cor(vector<float>* v1, vector<float>* v2){
	double toreturn;
	double m1 = 0;
	double m2 = 0;
	double sx =0;
	double ss1 = 0;
	double ss2 = 0;

	for (int i = 0; i < v1->size(); i++){
		m1 += v1->at(i);
		m2 += v2->at(i);
	}
	m1 = m1 / (double) v1->size();
	m2 = m2/ (double) v2->size();

	for (int i = 0;i < v1->size(); i++){
		sx += (v1->at(i) - m1)* (v2->at(i) - m2);
		ss1 += (v1->at(i) - m1)*(v1->at(i)- m1);
		ss2 += (v2->at(i) - m2)*(v2->at(i)- m2);
	}
	ss1 = sqrt(ss1);
	ss2 = sqrt(ss2);
	toreturn = sx / (ss1*ss2);
	return toreturn;
}

pair<double, double> HapMatrices::get_regress(vector<float>* v1, vector<float>* v2){
	pair<double, double> toreturn;
	double m1 = 0;
	double m2 = 0;
	double mx =0;
	double mss1 = 0;
	double mss2 = 0;

	for (int i = 0; i < v1->size(); i++){
		m1 += v1->at(i);
		m2 += v2->at(i);
		mx += v1->at(i)*v2->at(i);
		mss1 += v2->at(i)* v2->at(i);
	}
	//cout << m1 << " "<< m2 << " "<< mx << " "<< mss1 << " "<< v1->size() <<"\n";
	m1 = m1 / (double) v1->size();
	m2 = m2/ (double) v2->size();
	mx = mx / (double) v1->size();
	mss1 = mss1 / (double) v1->size();

	toreturn.second = (mx - m1*m2)/ (mss1 - m2*m2);
	toreturn.first = m1 - toreturn.second*m2;
	return toreturn;
}

pair<double, double> HapMatrices::get_regress_lowmem(vector<float> summary){
	pair<double, double> toreturn;
	double m1 = summary[0];
	double m2 = summary[1];
	double mx =summary[2];
	double mss1 = summary[3];
	double total = summary[4];
	//cout << "low "<< m1 << " "<< m2 << " "<< mx << " "<< mss1 << " "<< total << "\n";
	m1 = m1 / total;
	m2 = m2/ total;
	mx = mx / total;
	mss1 = mss1 / total;

	toreturn.second = (mx - m1*m2)/ (mss1 - m2*m2);
	toreturn.first = m1 - toreturn.second*m2;
	return toreturn;
}


vector<pair<double, double> > HapMatrices::rolloff(double by){
	vector<pair<double, double> > toreturn;
	cout << "Number of SNPs "<< set_weights();
	vector<pair<vector<float>, vector<float> > > binned_d = bin_d(by);
	for (double i  = 0; i < MAXDIST; i+= by){
		//cout << i << " ";
		double jtmp = i/by;
		int j = (int) jtmp;
		double mid = i+ by/2.0;
		pair<double, double> r = get_regress(&binned_d.at(j).first, &binned_d.at(j).second);
		toreturn.push_back(make_pair(mid, r.second));
	}
	return toreturn;
}

vector<pair<double, double> > HapMatrices::rolloff_lowmem(double by){
	vector<pair<double, double> > toreturn;
	nsnp = set_weights();
	cout << "Number of SNPs "<< nsnp <<"\n";
	vector<vector<float> > binned_d = bin_d_lowmem(by);
	//cout << "Binned_low\n"; cout.flush();
	for (double i  = 0; i < MAXDIST; i+= by){
		//cout << i << " ";
		double jtmp = i/by;
		int j = (int) jtmp;
		double mid = i+ by/2.0;
		pair<double, double> r = get_regress_lowmem(binned_d[j]);
		toreturn.push_back(make_pair(mid, r.second));
	}
	return toreturn;
}

vector<pair<double, double> > HapMatrices::rolloff_freq_lowmem(double by){
	vector<pair<double, double> > toreturn;
	nsnp = set_weights();
	cout << "Number of SNPs "<< nsnp <<"\n";
	vector<vector<float> > binned_d = bin_d_freq_lowmem(by);
	//cout << "Binned_low\n"; cout.flush();
	for (double i  = 0; i < MAXDIST; i+= by){
		//cout << i << " ";
		double jtmp = i/by;
		int j = (int) jtmp;
		double mid = i+ by/2.0;
		pair<double, double> r = get_regress_lowmem(binned_d[j]);
		toreturn.push_back(make_pair(mid, r.second));
	}
	return toreturn;
}


vector<pair<double, double> > HapMatrices::rolloff_freq_cor_lowmem(double by){
	vector<pair<double, double> > toreturn;
	nsnp = set_weights();
	cout << "Number of SNPs "<< nsnp <<"\n";
	vector<vector<float> > binned_d = bin_d_freq_lowmem(by);
	vector<pair<float, float> > cors;
	//cout << "Binned_low\n"; cout.flush();
	for (double i  = 0; i < MAXDIST; i+= by){
		//cout << i << " ";
		double jtmp = i/by;
		int j = (int) jtmp;
		double mid = i+ by/2.0;
		pair<double, double> r = get_regress_lowmem(binned_d[j]);
		cors.push_back(r);
		//toreturn.push_back(make_pair(mid, r.second));
	}
	cout << "Got frequencies\n";
	vector<vector<float> > binned_d2 = bin_d_freq_cor_lowmem(by, cors);
	for (double i  = 0; i < MAXDIST; i+= by){
		//cout << i << " ";
		double jtmp = i/by;
		int j = (int) jtmp;
		double mid = i+ by/2.0;
		pair<double, double> r = get_regress_lowmem(binned_d[j]);
		toreturn.push_back(make_pair(mid, r.second));
	}

	return toreturn;
}


pair<double, double> HapMatrices::fitexp(vector<pair<double, double> > d, int skip){
	pair<double, double> toreturn;
	vector<float> d1;
	vector<float> d2;
	for (int i = skip; i< d.size(); i++){
		d1.push_back(log(d[i].second));
		d2.push_back(d[i].first);
	}
	toreturn = get_regress(&d1, &d2);
	return toreturn;
}

void HapMatrices::roll_and_fit(string outfile, double by, int skip){
	ofstream out(outfile.c_str());
	vector<pair<double, double> > roll = rolloff(by);
	pair<double, double> fit = fitexp(roll, skip);
	for (vector<pair<double, double> >::iterator it = roll.begin(); it != roll.end(); it++){
		double x = it->first;
		double cor = it->second;
		double pred = exp(fit.first + fit.second*x);
		out << x << " "<< cor << " "<< pred << "\n";
	}
}

void HapMatrices::roll_print(string outfile, double by){
	ofstream out(outfile.c_str());
	vector<pair<double, double> > roll = rolloff(by);
	for (vector<pair<double, double> >::iterator it = roll.begin(); it != roll.end(); it++){
		out << it->first << " "<< it->second << "\n";
	}
}

void HapMatrices::roll_print_lowmem(string outfile, double by){
	ofstream out(outfile.c_str());
	vector<pair<double, double> > roll = rolloff_lowmem(by);
	//cout << beta_cor << "\n";
	out << "# beta_cor "<< beta_cor << "\n";

	for (vector<pair<double, double> >::iterator it = roll.begin(); it != roll.end(); it++){
		out << it->first << " "<< it->second*beta_cor << "\n";
	}
}

void HapMatrices::roll_print_freq_lowmem(string outfile, double by){
	ofstream out(outfile.c_str());
	vector<pair<double, double> > roll = rolloff_freq_lowmem(by);
	for (vector<pair<double, double> >::iterator it = roll.begin(); it != roll.end(); it++){
		out << it->first << " "<< it->second << "\n";
	}
}

void HapMatrices::roll_print_freq_cor_lowmem(string outfile, double by){
	ofstream out(outfile.c_str());
	vector<pair<double, double> > roll = rolloff_freq_cor_lowmem(by);
	for (vector<pair<double, double> >::iterator it = roll.begin(); it != roll.end(); it++){
		out << it->first << " "<< it->second << "\n";
	}
}
